import networkx as nx
import numpy as np
import sys
import os
import pylab as plt
import math
from params import *

datasets = ['bkite','fsquare','gowalla']
lbls = ['Brightkite', 'Foursquare', 'Gowalla']
node_dist = [5651.0,8494.0,5663.0]
link_dist = [2041.0,1442.0,1792.0]
degs = [7.88,22.07,9.48,7.0]
def log_fit(x,y):
    log_x = map(math.log,x)
    m1, m2 = min(log_x), max(log_x)
    x_fit = np.linspace(m1,m2,100)
    log_y = map(math.log,y)
    a,b = np.polyfit(log_x,log_y,1)
    avg_y = float(sum(log_y))/len(log_y)
    meanerror = 0.0
    residual = 0.0
    for x, y in zip(log_x,log_y):
        meanerror = meanerror + (y - avg_y)**2
        residual = residual + (y - a * x - b)**2
    RR = 1 - residual/meanerror
    print RR
    y_fit = [math.exp(a*i+b) for i in x_fit]
    x_fit = map(math.exp,x_fit)
    return x_fit, y_fit, a,b, RR


def scale_degree(d):
    M1 = 0.0
    M2 = 4.0
    MIN = 0
    MAX = 10**M2
    STEPS = 100
    l = math.log10(d)
    ratio = (l-M1)/(M2-M1)
    index = float(int(ratio*STEPS))/STEPS
    f = M1 + index*(M2-M1)
    #print d,l, ratio,index,f
    return 10**f

FILETYPE = 'pdf'
for dataset in datasets:
    print dataset

    def compute_distance_degree(graph_file):
        data = {}
        def update_node(n,d):
            if not n in data:
                data[n] = (0,0.0)
            deg,tot_dist = data[n]
            deg += 1
            tot_dist += d
            data[n] = (deg,tot_dist)

        avg_dist = 0.0
        k = 0
        for line in open(graph_file):
            k += 1
            n1,n2,d = line.strip().split()
            n1,n2 = map(int,(n1,n2))
            d = float(d)
            avg_dist += d
            update_node(n1,d)
            update_node(n2,d)

        avg_dist = avg_dist / k
        print "read %d nodes"%(len(data))
        print 'avg link length ', avg_dist

        values = {}
        for n in data:
            deg,tot_dist = data[n]
            deg = scale_degree(deg)
            values.setdefault(deg,[]).append(tot_dist)#/deg)

        print 'read %d degree values'%(len(values))

        def compute_avg(seq):
            m = sum(seq)/len(seq)
            m2 = sum(i**2 for i in seq)/len(seq)
            var = m2 - m**2
            std = var**0.5
            return m,std

        def compute_avg2(seq):
            s = map(math.log,seq)
            m = math.exp(sum(s)/len(s))
            #m2 = sum(i**2 for i in seq)/len(seq)
            #var = m2 - m**2
            #std = var**0.5
            return m,0

        x,y = [],[]
        for deg in sorted(values):
            x.append(deg)
            avg,std = compute_avg(values[deg])
            avg = avg/avg_dist
            y.append(avg)

        return x,y

    results = {}
    DIR = os.path.join(WORKDIR, 'gsn', 'null_graphs',dataset)
    for file in os.listdir(DIR):
        if not 'null_graph' in file:
            continue
        if 'er_null' in file:
            LBL = 'ER'
            continue
        elif 'geo_null' in file:
            LBL = 'GEO'
        elif 'social_null' in file:
            LBL = 'SOCIAL'
        else:
            continue
        num = int(file.split('.')[-2].split('_')[-1])
        if num != 2:
            continue

        print file, dataset, num

        graph_file = os.path.join(DIR,file)
        x,y = compute_distance_degree(graph_file)
        results[LBL] = (x,y)

    degree_file = os.path.join(WORKDIR, 'gsn','results',dataset, '%s_distance_degree.txt'%dataset)

    if 'bkite' in dataset:
        avg_dist = link_dist[0]
    elif 'fsquare' in dataset:
        avg_dist = link_dist[1]
    else:
        avg_dist = link_dist[2]

    degrees = []
    distances= []
    for line in open(degree_file):
        deg,dist = map(float,line.split(';'))
        degrees.append(deg)
        distances.append(deg*dist/avg_dist)
        #distances.append(dist/avg_dist)
    results['DATA'] = (degrees,distances)

    legs = []
    beta = []

    plt.figure()
    plt.clf()
    plt.axes(FIG_AXES2)
    keys = ('DATA','GEO','SOCIAL')
    lbls = ['Original data', 'Geo model', 'Social model']
    markers = ['x', '^', 'o']
    colors = list('rgb')

    for lbl,m1,m2,c in zip(keys,markers,symbols, colors):
        print lbl, m1, m2
        degrees,distances = results[lbl]
        ax1 = plt.loglog(degrees,distances,'%s%s'%(c,m1), mfc='None', mec=c)#,alpha=0.8)
        legs.append(ax1[0])
    for lbl,m1,m2 in zip(keys,markers,symbols):
        degrees,distances = results[lbl]
        x_fit,y_fit,a,b,RR = log_fit(degrees,distances)
        print '%s k = %f, A = %f'%(lbl,a,math.exp(b))
        ax1 = plt.loglog(x_fit,y_fit,'k--') #%s'%m2, mfc='None')#,linewidth=2)
        beta.append(a)
    #    plt.figtext(0.3,0.25,r'$s^d(k) = Ak^{\beta_{ER}} (\beta_{ER}=%.2f)$'%(
    #        a),size=20, bbox=dict(facecolor='w',fill=True,edgecolor='k'))
    #plt.title(r'$\beta_{ER} = %.2f, \beta_{GEO} = %.2f, \beta_{SOC} = %.2f$'%tuple(beta))
    #plt.title(r'$\beta_{GEO} = %.2f, \beta_{SOC} = %.2f$'%tuple(beta))
    plt.legend(legs,lbls,numpoints=1,
            loc='upper left')
    plt.grid(True)
    plt.xlabel(r'$k$')
    plt.ylabel('Average distance strength')
    plt.ylabel(r'$s(k)$')
    plt.xlim(xmax=1e3)
    plt.savefig('%s_null_distance_degree.%s'%(dataset,FILETYPE))
    plt.close()
